package com.kanq.ocr.service.impl;

import com.benjaminwan.ocrlibrary.OcrResult;
import com.kanq.ocr.config.OcrAiProperties;
import com.kanq.ocr.config.PromptTemplates;
import com.kanq.ocr.service.DocumentRecognitionService;
import com.kanq.ocr.service.OcrEngineManager;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.pdfbox.Loader;
import org.apache.pdfbox.io.IOUtils;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executor;
import java.util.stream.Collectors;

/**
 * @author wangchong
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class DocumentRecognitionServiceImpl implements DocumentRecognitionService {

	private final OcrAiProperties ocrAiProperties;
	private final OpenAiChatModel aiClient;
	private final PromptTemplates promptTemplates;
	@Qualifier("ocrTaskExecutor")
	private final Executor ocrTaskExecutor;


	@Override
	public String recognizeDocument(String templateType, List<?> inputs) {
		try {
			List<InputStream> inputStreams = convertInputs(inputs);
			TemplateType type = TemplateType.valueOf(templateType);
			switch (type) {
				case CONTRACT:
					return recognizeContract(inputStreams.get(0));
				case REAL_ESTATE:
					return recognizeRealEstate(inputStreams);
				case HOUSEHOLD_REGISTER:
					return recognizeHouseholdRegister(inputStreams);
				case ID_CARD:
					return recognizeIDCard(inputStreams);
				default:
					throw new IllegalArgumentException("不支持的模板类型: " + templateType);
			}
		} catch (IllegalArgumentException e) {
			throw new IllegalArgumentException("无效的模板类型: " + templateType, e);
		} catch (IOException e) {
			throw new RuntimeException(e);
		}
	}

	@Override
	public String recognizeContract(InputStream pdfInputStream) {
		try (PDDocument document = Loader.loadPDF(IOUtils.toByteArray(pdfInputStream))) {
			// 1. 提取PDF文本内容
			PDFTextStripper stripper = new PDFTextStripper();
			String rawText = stripper.getText(document);

			// 2. 清洗文本内容
			String cleanedText = cleanExtractedText(rawText);

			// 3. 提取关键信息生成JSON
			String initialJson = extractKeyInformation(cleanedText);

			// 4. 使用AI优化JSON格式
			String prompt = String.format(promptTemplates.getContractTemplate(), cleanedText, initialJson);
			return aiClient.call(prompt);

		} catch (Exception e) {
			log.error("合同处理流程异常", e);
			throw new RuntimeException("合同处理失败", e);
		}
	}

	// 新增文本清洗方法
	private String cleanExtractedText(String rawText) {
		// 移除多余空格、换行符和特殊字符
		return rawText.replaceAll("[\\u0000-\\u001F]", "")
				.replaceAll("\\s+", " ")
				.replaceAll("(?m)^\\s+$", "")
				.trim();
	}

	// 新增关键信息提取方法
	private String extractKeyInformation(String text) {
		return String.format("{\"sections\": [%s]}",
				text.lines()
						.filter(line -> ocrAiProperties.getKeywords().getContractFields()
								.stream().anyMatch(line::contains))
						.map(line -> String.format("\"%s\"", line.trim()))
						.collect(Collectors.joining(",")));
	}

	@Override
	public String recognizeRealEstate(List<InputStream> imageInputStreams) {
		try {
			// 1. 并行处理所有图片
			List<CompletableFuture<String>> futures = imageInputStreams.stream()
					.map(this::processImageAsync)
					.collect(Collectors.toList());

			// 2. 等待所有图片处理完成
			List<String> results = futures.stream()
					.map(CompletableFuture::join)
					.collect(Collectors.toList());

			// 3. 合并结果并使用AI处理
			String recognizedText = String.join("\n", results);
			String prompt = String.format(promptTemplates.getRealEstateTemplate(), recognizedText);
			return aiClient.call(prompt);

		} catch (Exception e) {
			log.error("不动产权证识别失败", e);
			throw new RuntimeException("不动产权证识别失败", e);
		}
	}

	@Override
	public String recognizeHouseholdRegister(List<InputStream> imageInputStreams) {
		try {
			// 1. 并行处理所有图片
			List<CompletableFuture<String>> futures = imageInputStreams.stream()
					.map(this::processImageAsync)
					.collect(Collectors.toList());

			// 2. 等待所有图片处理完成
			List<String> results = futures.stream()
					.map(CompletableFuture::join)
					.collect(Collectors.toList());

			// 3. 合并结果并使用AI处理
			String recognizedText = String.join("\n", results);
			String prompt = String.format(promptTemplates.getHouseholdRegisterTemplate(), recognizedText);
			return aiClient.call(prompt);

		} catch (Exception e) {
			log.error("户口本识别失败", e);
			throw new RuntimeException("户口本识别失败", e);
		}
	}

	@Override
	public String recognizeIDCard(List<InputStream> imageInputStreams) {
		try {
			// 1. 并行处理所有图片
			List<CompletableFuture<String>> futures = imageInputStreams.stream()
					.map(this::processImageAsync)
					.collect(Collectors.toList());

			// 2. 等待所有图片处理完成
			List<String> results = futures.stream()
					.map(CompletableFuture::join)
					.collect(Collectors.toList());

			// 3. 合并结果并使用AI处理
			String recognizedText = String.join("\n", results);
			String prompt = String.format(promptTemplates.getIdCardTemplate(), recognizedText);
			return aiClient.call(prompt);

		} catch (Exception e) {
			log.error("身份证识别失败", e);
			throw new RuntimeException("身份证识别失败", e);
		}
	}

	@Async("ocrTaskExecutor")
	public CompletableFuture<String> processImageAsync(InputStream imageInputStream) {
		try {
			BufferedImage image = ImageIO.read(imageInputStream);
			return CompletableFuture.completedFuture(processImage(image));
		} catch (Exception e) {
			log.error("图片处理失败", e);
			return CompletableFuture.failedFuture(e);
		}
	}

	private String processImage(BufferedImage image) {
		try {
			if (image == null) {
				throw new IllegalArgumentException("无法读取图片数据");
			}

			// 1. 将BufferedImage保存为临时文件
			File tempFile = File.createTempFile("ocr_", ".png");
			tempFile.deleteOnExit();
			ImageIO.write(image, "PNG", tempFile);

			// 2. 使用OCR引擎管理器执行识别
			OcrResult result = OcrEngineManager.performOcr(tempFile.getAbsolutePath());
			return result.getStrRes();
		} catch (Exception e) {
			log.error("OCR识别失败", e);
			throw new RuntimeException("OCR识别失败", e);
		}
	}

	private InputStream convertToInputStream(Object input) throws IOException {
		if (input instanceof MultipartFile) {
			return ((MultipartFile) input).getInputStream();
		} else if (input instanceof URL) {
			return ((URL) input).openStream();
		} else if (input instanceof File) {
			return new FileInputStream((File) input);
		} else if (input instanceof InputStream) {
			return (InputStream) input;
		} else {
			throw new IllegalArgumentException("不支持的输入类型: " + input.getClass().getName());
		}
	}

	private List<InputStream> convertInputs(List<?> inputs) throws IOException {
		List<InputStream> streams = new ArrayList<>();
		try {
			for (Object input : inputs) {
				streams.add(convertToInputStream(input));
			}
			return streams;
		} catch (Exception e) {
			// 确保已打开的流被关闭
			for (InputStream stream : streams) {
				stream.close();
			}
			throw e;
		}
	}
}